"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/7/2 9:20
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    parser_cam_param.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    parser_cam_param_VER0.1

------------------------------------------------------------------
# @Description:

------------------------------------------------------------------
"""
import json
import numpy as np


class ParserCamParam(object):
    def __init__(self, json_file_path="/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/"
                                      "rgb1920x1080_ir1280x800.json"):
        self.json_path = json_file_path
        self.json_info = None

    def load_cam_info(self):
        with open(self.json_path, "r") as f:
            data = json.load(f)

        self.json_info = data

    @staticmethod
    def cal_rotate_and_translation_coeffs(rotate_matrix, translation_matrix, color_intrinsics, depth_intrinsics,
                                          baseline=50.55):
        # Creating matrices and vectors based on given calibration data
        r_ir2rgb = np.array([
            [rotate_matrix['rot0'], rotate_matrix['rot1'], rotate_matrix['rot2']],
            [rotate_matrix['rot3'], rotate_matrix['rot4'], rotate_matrix['rot5']],
            [rotate_matrix['rot6'], rotate_matrix['rot7'], rotate_matrix['rot8']]
        ])

        k_rgb = np.array([
            [color_intrinsics['fx'], 0, color_intrinsics['cx']],
            [0, color_intrinsics['fy'], color_intrinsics['cy']],
            [0, 0, 1]
        ])

        k_ir = np.array([
            [depth_intrinsics['fx'], 0, depth_intrinsics['cx']],
            [0, depth_intrinsics['fy'], depth_intrinsics['cy']],
            [0, 0, 1]
        ])

        r_matrix = k_rgb.dot(r_ir2rgb).dot(np.linalg.inv(k_ir))

        t_ir = np.array([0, 0, 0])
        t_ir2rgb = np.array([
            translation_matrix['trans0'],
            translation_matrix['trans1'],
            translation_matrix['trans2']
        ])

        t_matrix = k_rgb.dot(r_ir2rgb.dot(t_ir) + t_ir2rgb)

        # Storing the results back into calib_data dictionary
        r_coeffs = {
            'r0': r_matrix[0, 0],
            'r1': r_matrix[0, 1],
            'r2': r_matrix[0, 2],
            'r3': r_matrix[1, 0],
            'r4': r_matrix[1, 1],
            'r5': r_matrix[1, 2],
            'r6': r_matrix[2, 0],
            'r7': r_matrix[2, 1],
            'r8': r_matrix[2, 2]
        }

        t_coeffs = {
            't0': t_matrix[0],
            't1': t_matrix[1],
            't2': t_matrix[2]
        }

        # Calculating Q matrix based on the baseline
        if baseline > 0:
            q_coeffs = {
                'Q23': depth_intrinsics['fx'],
                'Q32': 1 / baseline,
                'Q33': 0
            }
        else:
            print("Check input baseline value(> 0)!")
            q_coeffs = {depth_intrinsics['fx'], 0, 0}

        return r_coeffs, t_coeffs, q_coeffs

    def parser_cam_param(self, disp_img_w, disp_img_h):
        self.load_cam_info()

        cam_info = self.json_info["CameraModule"]["CameraInfo"]
        color_intrinsics = cam_info["ColorIntrinsic"]
        depth_intrinsics = cam_info["DepthIntrinsic"]
        d2c_transform_parm = cam_info["D2CTransformParam"]
        rotation_matrix = d2c_transform_parm["RotateMatrix"]
        translation_matrix = d2c_transform_parm["TranslateMatrix"]

        scale_w_ratio = 1.0
        scale_h_ratio = 1.0
        if disp_img_w != depth_intrinsics['width']:
            scale_w_ratio = disp_img_w / depth_intrinsics['width']

        if disp_img_h != depth_intrinsics['height']:
            scale_h_ratio = disp_img_h / depth_intrinsics['height']

        if scale_w_ratio != 1.0 or scale_h_ratio != 1.0:
            depth_intrinsics['width'] = disp_img_w
            depth_intrinsics['height'] = disp_img_h
            depth_intrinsics['fx'] *= scale_w_ratio
            depth_intrinsics['cx'] *= scale_w_ratio
            depth_intrinsics['fy'] *= scale_h_ratio
            depth_intrinsics['cy'] *= scale_h_ratio

        r_coeffs, t_coeffs, q_coeffs = ParserCamParam().cal_rotate_and_translation_coeffs(rotation_matrix,
                                                                                          translation_matrix,
                                                                                          color_intrinsics,
                                                                                          depth_intrinsics,
                                                                                          baseline=50.55)

        cam_param = {
            'color_intrinsics': color_intrinsics,
            'depth_intrinsics': depth_intrinsics,
            'r_coeffs': r_coeffs,
            't_coeffs': t_coeffs,
            'q_coeffs': q_coeffs
        }

        return cam_param, scale_w_ratio, scale_h_ratio


if __name__ == '__main__':
    parser = ParserCamParam()
    parser.parser_cam_param(640, 480)
